import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
confirmed = pd.read_csv('data/COVID/time_series_covid19_confirmed_global.csv').drop(columns=['Province/State', 'Lat', 'Long'])
deaths = pd.read_csv('data/COVID/time_series_covid19_deaths_global.csv').drop(columns=['Province/State', 'Lat', 'Long'])
recovered = pd.read_csv('data/COVID/time_series_covid19_recovered_global.csv').drop(columns=['Province/State', 'Lat', 'Long'])

cols = ['Country/Region']
confirmed = confirmed.groupby(cols).sum().reset_index()
deaths = deaths.groupby(cols).sum().reset_index()
recovered = recovered.groupby(cols).sum().reset_index()

confirmed = confirmed.melt(id_vars=cols, var_name='date', value_name='confirmed')
deaths = deaths.melt(id_vars=cols, var_name='date', value_name='deaths')
recovered = recovered.melt(id_vars=cols, var_name='date', value_name='recovered')
data = confirmed.merge(deaths, on=cols + ['date']).merge(recovered, on=cols + ['date'])
data['date'] = pd.to_datetime(data['date'], format='%m/%d/%y')

country_map = {
    'United States': 'US',
    'US': 'US',
    'United Kingdom': 'United Kingdom',
    'Brazil': 'Brazil',
    'India': 'India',
    'Russia': 'Russia',
    'South Africa': 'South Africa'
}
country_labels = {
    'US': '美国',
    'United Kingdom': '英国',
    'Brazil': '巴西',
    'India': '印度',
    'Russia': '俄罗斯',
    'South Africa': '南非'
}

countries = list(country_labels.keys())
focus = data[data['Country/Region'].isin(countries)].copy()
focus = focus.sort_values(['Country/Region', 'date'])
focus['daily_cases'] = focus.groupby('Country/Region')['confirmed'].diff().clip(lower=0)
focus['daily_deaths'] = focus.groupby('Country/Region')['deaths'].diff().clip(lower=0)
def seven_day_mean(series):
    return series.rolling(7, min_periods=1).mean()

daily_cases_ma = focus.groupby('Country/Region')['daily_cases'].apply(seven_day_mean)
daily_cases_ma = daily_cases_ma.reset_index(level=0, drop=True)
focus['daily_cases_ma'] = daily_cases_ma

def shift_five(series):
    return series.shift(5)

rt_series = focus.groupby('Country/Region')['daily_cases_ma'].apply(shift_five)
rt_series = focus['daily_cases_ma'] / rt_series.reset_index(level=0, drop=True)
focus['Rt'] = rt_series
focus['Rt'] = focus['Rt'].replace([np.inf, -np.inf], np.nan)

plt.figure(figsize=(10, 5))
for key in countries:
    part = focus[focus['Country/Region'] == key]
    plt.plot(part['date'], part['Rt'], label=country_labels.get(key, key))
plt.axhline(1, color='black', linestyle='--')
plt.ylim(0, 4)
plt.xlabel('日期')
plt.ylabel('Rt')
plt.legend()
plt.tight_layout()
plt.savefig('第二次平时作业_第一题_图1.png', dpi=300)
plt.close()

usa = focus[focus['Country/Region'] == 'US'].copy()
usa['cases_ma'] = usa['daily_cases'].rolling(7, min_periods=1).mean()
usa['deaths_ma'] = usa['daily_deaths'].rolling(7, min_periods=1).mean()
plt.figure(figsize=(10, 4))
plt.plot(usa['date'], usa['cases_ma'], label='病例')
plt.plot(usa['date'], usa['deaths_ma'], label='死亡')
plt.xlabel('日期')
plt.legend()
plt.tight_layout()
plt.savefig('第二次平时作业_第一题_图2.png', dpi=300)
plt.close()

def shift_fourteen(series):
    return series.shift(14)

cases_shift = focus.groupby('Country/Region')['daily_cases_ma'].apply(shift_fourteen)
focus['cases_shift'] = cases_shift.reset_index(level=0, drop=True)
focus['death_rate'] = focus['daily_deaths'] / focus['cases_shift']
plt.figure(figsize=(10, 5))
for key in countries:
    part = focus[focus['Country/Region'] == key]
    plt.plot(part['date'], part['death_rate'].clip(upper=0.1), label=country_labels.get(key, key))
plt.xlabel('日期')
plt.ylabel('死亡率')
plt.legend()
plt.tight_layout()
plt.savefig('第二次平时作业_第一题_图3.png', dpi=300)
plt.close()
